#!/usr/bin/env python
import os,sys,time

# add self to search path for testing
if __name__ == '__main__':
    sys.path.append("../python_modules")

import general_util
from wrap_text import wrap_text

import re
class fasta_sequences (object):
    
    def __init__ (self, identifier_regex = ">\s*([^\s]+)\s*", comment_header = '#'):
        """
        __init__

        """
        self.sequences_from_name = dict()
        self.name_from_index = list()
        self.id_regex = re.compile(identifier_regex)
        self.comment_header = set(comment_header)

    def count (self):
        return len(self.name_from_index)

    def all_sequences (self):
        for x in range(self.count()):
            seq_id = self.name_from_index[x]
            yield [seq_id, self.sequences_from_name[seq_id]];

    #---------------------------------------------------------------------------
    def __getitem__( self, name ):
        try:
            return self.sequences_from_name[name]
        except KeyError:
            return self.sequences_from_name[self.name_from_index[name] ]


    #---------------------------------------------------------------------------
    def has_key( self, name ):
        return self.sequences_from_name.has_key(name)

    #---------------------------------------------------------------------------
    #
    #   read from fasta sequence
    # 
    def read_dna_from_file (self, file_name):
        """
        read_dna_from_file

            opens file and calls read dna
            
        """
        dna_file = open(file_name)
        self.read_dna(dna_file)

    
    #---------------------------------------------------------------------------
    #
    #   read from fasta sequence
    # 
    def read_dna (self, lines):
        """
        read_dna

            assumes file already open

            file contains either raw dna or in fasta format
            
        """
        
        
        curr_seq = ''
        curr_id = "RAW"

        line_count = 0
        for line in lines:

            line_count += 1
            

            # remove new line
            line = line.rstrip('\n')

            if not len(line):
                continue

            # ignore comment lines
            if line[0] in self.comment_header:
                continue
            
            # fasta header lines
            if line[0] == '>':
                if len(curr_seq):
                    self.sequences_from_name[curr_id] = curr_seq
                    self.name_from_index.append(curr_id)
                elif curr_id != 'RAW':
                    raise "Empty sequence on line #%d " % ( line_count - 1)
                    
                curr_seq = ''
    
                # match identifier using regex, else use whole line
                id_match = self.id_regex.match(line)
                if id_match and id_match.group(1):
                    curr_id = id_match.group(1)
                else:
                    curr_id = line[1:]
                if len(curr_id) == 0:
                    raise "No identifier extracted on line # %d : [%s]" % (
                                line_count, line)

                # check no duplicate identifiers
                if curr_id in self.sequences_from_name:
                    raise "Duplicate identifiers [%s] on line # %d" % (
                                curr_id, line_count)
    
            # append sequence without new line
            else:
                curr_seq += line
                
        # save last sequence
        else:
            if len(curr_seq):
                self.sequences_from_name[curr_id] = curr_seq
                self.name_from_index.append(curr_id)


        
def unit_test():
    fa = fasta_sequences()
    seq_lines = \
"""
>abc sdf
011111111110
233333333333334
55
>efg
ABBBBBBBBBBBBBBA
CDDDDDDDDDDDC
55
>hij
EF
"""
    fa.read_dna(seq_lines.split('\n'))
    for seq_id, seq in fa.all_sequences():
        print ">"+ seq_id, "\n", seq

if __name__ == '__main__':
    unit_test()    
